import asyncio
from typing import Dict, Any, Optional
import redis
import json
import uuid
from datetime import datetime, timedelta
import logging
from app.core.config import settings

logger = logging.getLogger(__name__)


class TaskQueue:
    """Redis-based task queue for managing colorization jobs"""
    
    def __init__(self):
        self.redis_client = redis.from_url(settings.redis_url)
        self.task_prefix = "colorize:task:"
        self.queue_key = "colorize:queue"
        
    async def add_task(self, task_data: Dict[str, Any]) -> str:
        """Add a new task to the queue"""
        task_id = str(uuid.uuid4())
        task_data.update({
            "task_id": task_id,
            "created_at": datetime.utcnow().isoformat(),
            "status": "pending"
        })
        
        # Store task data
        self.redis_client.setex(
            f"{self.task_prefix}{task_id}",
            timedelta(hours=24),
            json.dumps(task_data)
        )
        
        # Add to queue
        self.redis_client.lpush(self.queue_key, task_id)
        
        logger.info(f"Task {task_id} added to queue")
        return task_id
    
    async def get_task(self, timeout: int = 30) -> Optional[Dict[str, Any]]:
        """Get the next task from the queue"""
        try:
            # Block and pop from queue
            result = self.redis_client.brpop(self.queue_key, timeout=timeout)
            if not result:
                return None
            
            _, task_id = result
            
            # Get task data
            task_data = self.redis_client.get(f"{self.task_prefix}{task_id}")
            if not task_data:
                return None
            
            return json.loads(task_data)
            
        except Exception as e:
            logger.error(f"Error getting task from queue: {e}")
            return None
    
    async def update_task_status(
        self,
        task_id: str,
        status: str,
        progress: int = 0,
        result: Optional[Dict[str, Any]] = None,
        error: Optional[str] = None
    ):
        """Update task status"""
        try:
            task_key = f"{self.task_prefix}{task_id}"
            task_data = self.redis_client.get(task_key)
            
            if not task_data:
                logger.warning(f"Task {task_id} not found")
                return
            
            task = json.loads(task_data)
            task.update({
                "status": status,
                "progress": progress,
                "updated_at": datetime.utcnow().isoformat()
            })
            
            if result:
                task["result"] = result
            if error:
                task["error"] = error
            
            # Update TTL
            self.redis_client.setex(
                task_key,
                timedelta(hours=24),
                json.dumps(task)
            )
            
        except Exception as e:
            logger.error(f"Error updating task status: {e}")
    
    async def get_task_status(self, task_id: str) -> Optional[Dict[str, Any]]:
        """Get task status and details"""
        try:
            task_data = self.redis_client.get(f"{self.task_prefix}{task_id}")
            if not task_data:
                return None
            
            return json.loads(task_data)
            
        except Exception as e:
            logger.error(f"Error getting task status: {e}")
            return None
    
    async def get_queue_length(self) -> int:
        """Get the current queue length"""
        try:
            return self.redis_client.llen(self.queue_key)
        except Exception as e:
            logger.error(f"Error getting queue length: {e}")
            return 0
    
    async def cleanup_old_tasks(self, max_age_hours: int = 24):
        """Clean up old completed tasks"""
        try:
            pattern = f"{self.task_prefix}*"
            keys = self.redis_client.keys(pattern)
            
            cutoff_time = datetime.utcnow() - timedelta(hours=max_age_hours)
            
            for key in keys:
                task_data = self.redis_client.get(key)
                if task_data:
                    task = json.loads(task_data)
                    created_at = datetime.fromisoformat(task["created_at"])
                    
                    if created_at < cutoff_time:
                        self.redis_client.delete(key)
                        
            logger.info("Cleaned up old tasks")
            
        except Exception as e:
            logger.error(f"Error cleaning up old tasks: {e}")
    
    async def get_queue_stats(self) -> Dict[str, Any]:
        """Get queue statistics"""
        try:
            queue_length = await self.get_queue_length()
            
            # Count tasks by status
            pattern = f"{self.task_prefix}*"
            keys = self.redis_client.keys(pattern)
            
            stats = {
                "pending": 0,
                "processing": 0,
                "completed": 0,
                "failed": 0,
                "total": len(keys),
                "queue_length": queue_length
            }
            
            for key in keys:
                task_data = self.redis_client.get(key)
                if task_data:
                    task = json.loads(task_data)
                    status = task.get("status", "unknown")
                    if status in stats:
                        stats[status] += 1
            
            return stats
            
        except Exception as e:
            logger.error(f"Error getting queue stats: {e}")
            return {"error": str(e)}


# Global queue instance
task_queue = TaskQueue()
